#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
agent controller
"""

import tornado
import json
import asyncio
from pydantic import ValidationError

from src.common.error import Code, Error
from src.common.http_response import Response, resp_ok, resp_err
from src.common.context import ctx_http_resp
from src.common.log import log_http_context, logger
from src.controller.executor import Executor
import src.resource.service as service


class AgentHandler(tornado.web.RequestHandler):
    """
    AgentHandler
    """
    executor = Executor()

    @tornado.gen.coroutine
    @log_http_context
    def post(self):
        """
        post 请求
        """
        req = tornado.escape.json_decode(self.request.body)
        stream = req.get('stream', False)
        if stream:
            yield self.stream_response(req)
        else:
            res = yield self.common_response(req)
            resp = res.model_dump()
            ctx_http_resp.set(resp)
            self.finish(resp)
    
    @tornado.concurrent.run_on_executor
    @log_http_context
    def stream_response(self, request):
        res = self.call_agent(request)
        self.set_header('Content-Type', 'text/event-stream')
        self.set_header('Cache-Control', 'no-cache')
        self.set_header('Connection', 'keep-alive')
        self._auto_finish = False
        try:
            loop = asyncio.get_event_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
        for item in res:
            if self.request.connection.stream.closed():
                break
            self.write(f"data: {item.model_dump_json()}\n\n")
            self.flush()
        self.finish()
        return

    @tornado.concurrent.run_on_executor
    @log_http_context
    def common_response(self, request):
        """
        run 业务逻辑
        """
        return self.call_agent(request)
    
    def call_agent(self, request):
        try:
            if isinstance(request, str):
                params = tornado.escape.json_decode(request)
            elif isinstance(request, dict):
                params = request
            else:
                err = Error(Code.BAD_REQUEST)
                return resp_err(err)
            # service
            response = service.agent_service.run(params)

        except (ValidationError, Error, Exception) as e:
            return resp_err(e)
        return response